'''
Created on Jul 8, 2014

@author: Siyuan Qi
'''

# Dependencies
import pymc as pm
import numpy as np

# Local imports
from sputil import smallNumber

# This is a user-defined function to specify the distribution 
# parameters of the attributes of entities
class DistParam:
	def __init__(self, element):
		def __toFloat(strList):
			if type(strList) is list:
				return [__toFloat(s) for s in strList]
			else:
				return float(strList)

		self.xMu = [float(mu) for mu in element.find('x').get('mu').split()]
		self.xVar = [float(var) for var in element.find('x').get('var').split()]
		self.xPosProbs = [float(p) for p in element.find('x').get('prob').split()]
		self.yMu = [float(mu) for mu in element.find('y').get('mu').split()]
		self.yVar = [float(var) for var in element.find('y').get('var').split()]
		self.yPosProbs = [float(p) for p in element.find('y').get('prob').split()]
		self.zMu = [float(mu) for mu in element.find('z').get('mu').split()]
		self.zVar = [float(var) for var in element.find('z').get('var').split()]
		self.zPosProbs = [float(p) for p in element.find('z').get('prob').split()]

		self.sizeMu = __toFloat([mu.split() for mu in element.find('size').get('mu').split(';')])
		self.sizeVar = __toFloat([var.split() for var in element.find('size').get('var').split(';')])
		self.sizeProbs = [float(p) for p in element.find('size').get('prob').split()]

		self.thetaMu = [float(mu)*np.pi/180 for mu in element.find('theta').get('mu').split()]
		self.thetaVar = [float(var)*np.pi/180 for var in element.find('theta').get('var').split()]
		self.thetaProbs = [float(p) for p in element.find('theta').get('prob').split()]
		
	
	def updateVar(self, ratio):
		def __update(varList, ratio):
			if type(varList) is list:
				return [__update(v, ratio) for v in varList]
			else:
				return varList*ratio
			
		self.xVar = [v*ratio for v in self.xVar]
		self.yVar = [v*ratio for v in self.yVar]
# 		self.zVar = [v*ratio for v in self.zVar]
# 		self.sizeVar = [update(v, ratio) for v in self.sizeVar]
# 		self.thetaVar = [update(v, ratio) for v in self.thetaVar]
		
		

# This is a user-defined function to specify the structure of an entity
class Entity:
	def __init__(self, node, parent):
		# Initialize
		self.node = node
		self.parent = parent
		self.children = None

		# Position
		# xPos, yPos, zPos are the mean positions of the mixture Gaussian model 
		# adding the parent's sizes and the alignment offset
		self.xPos = [0, 0, 0]
		self.yPos = [0, 0, 0]
		self.zPos = [0, 0, 0]
		self.x = 0
		self.y = 0
		self.z = 0

		self.aCoor = [0, 0, 1]

		# Size, follows a three-dimensional Gaussian distribution
		self.size = [0, 0, 0]

		# Orientation
		self.thetaBase = 0
		self.theta = 0
		self.thetaWorld = 0

		# Transformation matrices
		self.transToParent = np.matrix([[1, 0, 0], \
						  				[0, 1, 0], \
						  				[0, 0, 1]])
		self.transToWorld = np.matrix([[1, 0, 0], \
						  			   [0, 1, 0], \
						  			   [0, 0, 1]])

		# Probabilities
		self.prior = 0

	def sampleAttribute(self):
		""" Sample the attributes of current entity.
		
		Sample the attributes of the current entity.
		Notice that this function does not sample children of the current entity.
		It only samples the geometry attributes using top-down approach, and it
		does not affect the tree structure.
		
		"""
		def __baseOrientation():
			if self.parent:
				distIdx = np.argsort([abs(self.y), abs(self.x), abs(self.parent.size[1] - self.y), abs(self.parent.size[0] - self.x)])
				self.thetaBase = distIdx[0]*np.pi/2
			else:
				self.thetaBase = 0
			return self.thetaBase

		def __calMeanPosition():
			if self.parent:
				for i in range(3):
					self.xPos[i] = 0.5 * i * self.parent.size[0] + d.xMu[i]
					self.yPos[i] = 0.5 * i * self.parent.size[1] + d.yMu[i]
					self.zPos[i] = 0.5 * i * self.parent.size[2] + d.zMu[i]
			else:
				for i in range(3):
					self.xPos[i] = d.xMu[i]
					self.yPos[i] = d.yMu[i]
					self.zPos[i] = d.zMu[i]

		d = self.node.distribution

		__calMeanPosition()

		# Sample the alignment branch, value chosen from [0, 0.5, 1]
		xPosBranch = pm.distributions.rcategorical(d.xPosProbs)
		yPosBranch = pm.distributions.rcategorical(d.yPosProbs)
		zPosBranch = pm.distributions.rcategorical(d.zPosProbs)

		# Add a noise to the mean position
		self.x = pm.distributions.rnormal(self.xPos[xPosBranch], 1./d.xVar[xPosBranch])
		self.y = pm.distributions.rnormal(self.yPos[yPosBranch], 1./d.yVar[yPosBranch])
		self.z = pm.distributions.rnormal(self.zPos[zPosBranch], 1./d.zVar[zPosBranch])

		# Sample the size of the entity
		sizeBranch = pm.distributions.rcategorical(d.sizeProbs)
		self.size = pm.distributions.rmv_normal_cov(d.sizeMu[sizeBranch], np.diag(d.sizeVar[sizeBranch]))

		# Sample the orientation of the entity
		thetaBranch = pm.distributions.rcategorical(d.thetaProbs)
		self.theta = __baseOrientation() + pm.distributions.rnormal(d.thetaMu[thetaBranch], 1/d.thetaVar[thetaBranch])
		
		self._calParentTrans()
		self._calWorldTrans()


	def assignCuboid(self, cuboid):
		""" Assign a cuboid to the current entity.
		
		Assign a cuboid in the world coordinate system to the current entity and keep all the 
		other cuboids at their place.  Hence the relationship between the current entity and its 
		parent and the children entities need to be modified.
		Notice that each entity keeps its own relation with its parent entity.  Here we need to
		modify the parameters in the current entity and the children entities.
		First, the following attributes of the current entity will be modified:
			size, transToWorld, thetaWorld, theta, transToParent, x, y, z, thetaBase
		Then the following attributes of the children entities will be modified:
			xPos, yPos, zPos, theta, transToParent, x, y, thetaBase
		
		Args:
			cuboid: a list of length 6: [x1, y1, z1, x2, y2, z2]. Here the coordinates are based
				on the world coordinate system. (x1, y1, z1) is the coordinate of the origin
				of the cuboid (which determines the orientation of the cuboid).
		
		"""
		# First check the orientation of the cuboid to find out the correnspondence of the 
		# three dimensions of size to the dimensions of the cuboid
		if (cuboid[3] - cuboid[0])*(cuboid[4] - cuboid[1]) > 0:  
			# Calculate the new size of the cuboid
			self.size[0] = abs(cuboid[3] - cuboid[0])  # length corresponds to x
			self.size[1] = abs(cuboid[4] - cuboid[1])  # width corresponds to y
			self.size[2] = abs(cuboid[5] - cuboid[2])
			self.z = cuboid[2]

			# Calculate the new world coordinate of the achor point(rotation center) of the cuboid.
			# Notice that anchor[0] and anchor[1] corresponds to length and width respectively.
			# cuboid[0] and cuboid[1] corresponds to x and y repectively.
			wCoor = np.matrix([[cuboid[0] + self.node.anchor[0]*(cuboid[3] - cuboid[0])], \
							   [cuboid[1] + self.node.anchor[1]*(cuboid[4] - cuboid[1])], \
							   [1]])

			# Calculate the new transformation matrix from local to world.
			# Use three points to calculate the transformation.
			# Since length corresponds to x, the second point changes in x dimension.
			cubWolrd = np.matrix([[cuboid[0], cuboid[1], 1], \
								  [cuboid[3], cuboid[1], 1], \
								  [cuboid[3], cuboid[4], 1]]).T
			cubLocal = self._getCorners()[:, 0:3]
			self.transToWorld = cubWolrd*np.linalg.inv(cubLocal)
		else:
			# Calculate the new size of the cuboid
			self.size[0] = abs(cuboid[4] - cuboid[1])  # length corresponds to y
			self.size[1] = abs(cuboid[3] - cuboid[0])  # width corresponds to x
			self.size[2] = abs(cuboid[5] - cuboid[2])
			self.z = cuboid[2]

			# Calculate the new world coordinate of the achor point(rotation center) of the cuboid.
			# Notice that anchor[0] and anchor[1] corresponds to length and width respectively.
			# cuboid[0] and cuboid[1] corresponds to x and y repectively.
			wCoor = np.matrix([[cuboid[0] + self.node.anchor[1]*(cuboid[3] - cuboid[0])], \
							   [cuboid[1] + self.node.anchor[0]*(cuboid[4] - cuboid[1])], \
							   [1]])

			# Calculate the new transformation matrix from local to world
			# Use three points to calculate the transformation.
			# Since length corresponds to y, the second point changes in y dimension.
			cubWolrd = np.matrix([[cuboid[0], cuboid[1], 1], \
								  [cuboid[0], cuboid[4], 1], \
								  [cuboid[3], cuboid[4], 1]]).T
			cubLocal = self._getCorners()[:, 0:3]
			self.transToWorld = cubWolrd*np.linalg.inv(cubLocal)

		# Calculate the orientation to world
		if self.transToWorld[0, 0] > 1:  # The cosine value of theta might overflow due to precision issues
			cosVal = 1
		elif self.transToWorld[0, 0] < -1:
			cosVal = -1
		else:
			cosVal = self.transToWorld[0, 0]
		self.thetaWorld = np.arccos(cosVal)
		if self.transToWorld[0, 1] < 0:  # If the sin value of thetaWorld is negative, flip the sign
			self.thetaWorld = -self.thetaWorld

		# Calculate the new transformation matrix from local to parent
		if self.parent:
			self.transToParent = np.linalg.inv(self.parent.transToWorld) * self.transToWorld
			self.theta = self.thetaWorld - self.parent.thetaWorld
		else:
			self.transToParent = self.transToWorld
			self.theta = self.thetaWorld

		# Calculate the coordinate of anchor point in parent
		lCoor = self.transToParent * np.linalg.inv(self.transToWorld) * wCoor
		self.x = lCoor[0, 0]
		self.y = lCoor[1, 0]

		# Calculate the base orientation of the cuboid
		if self.parent:
			distIdx = np.argsort([abs(self.y), abs(self.x), abs(self.parent.size[1] - self.y), abs(self.parent.size[0] - self.x)])
			self.thetaBase = distIdx[0]*np.pi/2
		else:
			self.thetaBase = self.theta
			
		# Recalculate the rotation part of transformation matrices, to avoid precision problems (invalid cos and sin values)
		self.transToWorld[0:2, 0:2] = np.matrix([[np.cos(self.thetaWorld), np.sin(self.thetaWorld)], \
						  			   			 [-np.sin(self.thetaWorld), np.cos(self.thetaWorld)]])
		self.transToParent[0:2, 0:2] = np.matrix([[np.cos(self.theta), np.sin(self.theta)], \
						  			   			  [-np.sin(self.theta), np.cos(self.theta)]])

		# Modify the attributes of children
		if self.children:
			for child in self.children:
				# Calculate the mean positions of children in the new cuboid
				for i in range(0, 3):
					child.xPos[i] = 0.5 * i * self.size[0] + child.node.distribution.xMu[i]
					child.yPos[i] = 0.5 * i * self.size[1] + child.node.distribution.yMu[i]
					child.zPos[i] = 0.5 * i * self.size[2] + child.node.distribution.zMu[i]

				# Calculate the orientation of child to parent and the transformation matrix from child to parent
				child.theta = child.thetaWorld - self.thetaWorld
				child.transToParent = np.linalg.inv(self.transToWorld) * child.transToWorld
				# Recalculate the rotation to avoid overflow caused by precision problems
				child.transToParent[0:2, 0:2] = np.matrix([[np.cos(child.theta), np.sin(child.theta)], \
						  			   					   [-np.sin(child.theta), np.cos(child.theta)]])
				
				# Calculate the coordinate of child anchor in parent
				aCoorParent = child.transToParent * np.matrix([[child.node.anchor[0]*child.size[0]], [child.node.anchor[1]*child.size[1]], [1]])
				child.x = aCoorParent[0, 0]
				child.y = aCoorParent[1, 0]

				# Calculate the base orientation by calculating the distance to each wall
				distIdx = np.argsort([abs(child.y), abs(child.x), abs(self.size[1] - child.y), abs(self.size[0] - child.x)])
				child.thetaBase = distIdx[0]*np.pi/2

	def assignGroupCuboid(self, cuboid, child = None):
		""" Assign a cuboid to the current entity and keep the relative positions between the
			current entity and the children entities.
		
		Assign a cuboid in the world coordinate system to the current entity and keep all the 
		relationships between this entity and its children.  Hence the attributes of the children
		are kept as original.
		The following attributes of the current entity will be modified:
			size, transToWorld, thetaWorld, theta, transToParent, x, y, z, thetaBase
		
		Args:
			cuboid: a list of length 6: [x1, y1, z1, x2, y2, z2]. Here the coordinates are based
				on the world coordinate system. (x1, y1, z1) is the coordinate of the origin
				of the cuboid (which determines the orientation of the cuboid).
		
		"""
		# First check the orientation of the cuboid to find out the correnspondence of the 
		# three dimensions of size to the dimensions of the cuboid
		if (cuboid[3] - cuboid[0])*(cuboid[4] - cuboid[1]) > 0:  
			# Calculate the new size of the cuboid
			self.size[0] = abs(cuboid[3] - cuboid[0])  # length corresponds to x
			self.size[1] = abs(cuboid[4] - cuboid[1])  # width corresponds to y
			self.size[2] = abs(cuboid[5] - cuboid[2])
			self.z = cuboid[2]

			# Calculate the new world coordinate of the achor point(rotation center) of the cuboid.
			# Notice that anchor[0] and anchor[1] corresponds to length and width respectively.
			# cuboid[0] and cuboid[1] corresponds to x and y repectively.
			wCoor = np.matrix([[cuboid[0] + self.node.anchor[0]*(cuboid[3] - cuboid[0])], \
							   [cuboid[1] + self.node.anchor[1]*(cuboid[4] - cuboid[1])], \
							   [1]])

			# Calculate the new transformation matrix from local to world.
			# Use three points to calculate the transformation.
			# Since length corresponds to x, the second point changes in x dimension.
			cubWolrd = np.matrix([[cuboid[0], cuboid[1], 1], \
								  [cuboid[3], cuboid[1], 1], \
								  [cuboid[3], cuboid[4], 1]]).T
			cubLocal = self._getCorners()[:, 0:3]
			self.transToWorld = cubWolrd*np.linalg.inv(cubLocal)
		else:
			# Calculate the new size of the cuboid
			self.size[0] = abs(cuboid[4] - cuboid[1])  # length corresponds to y
			self.size[1] = abs(cuboid[3] - cuboid[0])  # width corresponds to x
			self.size[2] = abs(cuboid[5] - cuboid[2])
			self.z = cuboid[2]

			# Calculate the new world coordinate of the achor point(rotation center) of the cuboid.
			# Notice that anchor[0] and anchor[1] corresponds to length and width respectively.
			# cuboid[0] and cuboid[1] corresponds to x and y repectively.
			wCoor = np.matrix([[cuboid[0] + self.node.anchor[1]*(cuboid[3] - cuboid[0])], \
							   [cuboid[1] + self.node.anchor[0]*(cuboid[4] - cuboid[1])], \
							   [1]])

			# Calculate the new transformation matrix from local to world
			# Use three points to calculate the transformation.
			# Since length corresponds to y, the second point changes in y dimension.
			cubWolrd = np.matrix([[cuboid[0], cuboid[1], 1], \
								  [cuboid[0], cuboid[4], 1], \
								  [cuboid[3], cuboid[4], 1]]).T
			cubLocal = self._getCorners()[:, 0:3]
			self.transToWorld = cubWolrd*np.linalg.inv(cubLocal)

		# Calculate the orientation to world
		if self.transToWorld[0, 0] > 1:  # The cosine value of theta might overflow due to precision issues
			cosVal = 1
		elif self.transToWorld[0, 0] < -1:
			cosVal = -1
		else:
			cosVal = self.transToWorld[0, 0]
		self.thetaWorld = np.arccos(cosVal)
		if self.transToWorld[0, 1] < 0:  # If the sin value of thetaWorld is negative, flip the sign
			self.thetaWorld = -self.thetaWorld

		# Calculate the new transformation matrix from local to parent
		if self.parent:
			self.transToParent = np.linalg.inv(self.parent.transToWorld) * self.transToWorld
			self.theta = self.thetaWorld - self.parent.thetaWorld
		else:
			self.transToParent = self.transToWorld
			self.theta = self.thetaWorld

		# Calculate the coordinate of anchor point in parent
		lCoor = self.transToParent * np.linalg.inv(self.transToWorld) * wCoor
		self.x = lCoor[0, 0]
		self.y = lCoor[1, 0]

		# Calculate the base orientation of the cuboid
		if self.parent:
			distIdx = np.argsort([abs(self.y), abs(self.x), abs(self.parent.size[1] - self.y), abs(self.parent.size[0] - self.x)])
			self.thetaBase = distIdx[0]*np.pi/2
		else:
			self.thetaBase = self.theta
			
		# Recalculate the rotation part of transformation matrices, to avoid precision problems (invalid cos and sin values)
		self.transToWorld[0:2, 0:2] = np.matrix([[np.cos(self.thetaWorld), np.sin(self.thetaWorld)], \
						  			   			 [-np.sin(self.thetaWorld), np.cos(self.thetaWorld)]])
		self.transToParent[0:2, 0:2] = np.matrix([[np.cos(self.theta), np.sin(self.theta)], \
						  			   			  [-np.sin(self.theta), np.cos(self.theta)]])

		# Modify the attributes of child
		if child:
			# Calculate the mean positions of children in the new cuboid
			for i in range(0, 3):
				child.xPos[i] = 0.5 * i * self.size[0] + child.node.distribution.xMu[i]
				child.yPos[i] = 0.5 * i * self.size[1] + child.node.distribution.yMu[i]
				child.zPos[i] = 0.5 * i * self.size[2] + child.node.distribution.zMu[i]

			# Calculate the orientation of child to parent and the transformation matrix from child to parent
			child.theta = child.thetaWorld - self.thetaWorld
			child.transToParent = np.linalg.inv(self.transToWorld) * child.transToWorld
			# Recalculate the rotation to avoid overflow caused by precision problems
			child.transToParent[0:2, 0:2] = np.matrix([[np.cos(child.theta), np.sin(child.theta)], \
					  			   					   [-np.sin(child.theta), np.cos(child.theta)]])
			
			# Calculate the coordinate of child anchor in parent
			aCoorParent = child.transToParent * np.matrix([[child.node.anchor[0]*child.size[0]], [child.node.anchor[1]*child.size[1]], [1]])
			child.x = aCoorParent[0, 0]
			child.y = aCoorParent[1, 0]

			# Calculate the base orientation by calculating the distance to each wall
			distIdx = np.argsort([abs(child.y), abs(child.x), abs(self.size[1] - child.y), abs(self.size[0] - child.x)])
			child.thetaBase = distIdx[0]*np.pi/2


	def _calParentTrans(self):
		H1 = np.matrix([[1, 0, -self.node.anchor[0]*self.size[0]], [0, 1, -self.node.anchor[1]*self.size[1]], [0, 0, 1]])

		H2 = np.matrix([[np.cos(self.theta), np.sin(self.theta), self.x], \
					  	[-np.sin(self.theta), np.cos(self.theta), self.y], \
					  	[0, 0, 1]])

		self.transToParent = H2*H1

	def _calWorldTrans(self):
		if self.parent:
			self.thetaWorld = self.parent.thetaWorld + self.theta
			self.transToWorld = self.parent.transToWorld * self.transToParent
		else:
			self.thetaWorld = self.theta
			self.transToWorld = np.matrix([[np.cos(self.thetaWorld), np.sin(self.thetaWorld), self.x], \
							  			   [-np.sin(self.thetaWorld), np.cos(self.thetaWorld), self.y], \
							  			   [0, 0, 1]])

	def _getCorners(self):
		scale = np.matrix([[self.size[0], 0, 0], \
									 [0, self.size[1], 0], \
									 [0, 0, 1]])

		corners = np.matrix([[0, 0, 1], \
							 [1, 0, 1], \
							 [1, 1, 1], \
							 [0, 1, 1]]).T

		return scale*corners

	def _getParentCoor(self):
		return self.transToParent*self._getCorners()

	def _getWorldCoor(self):
		self._calWorldTrans()
		return self.transToWorld*self._getCorners()
	
	def getCuboid(self):
		coor = self._getWorldCoor()
		cuboid = [coor[0, 0], coor[1, 0], self.z, coor[0, 2], coor[1, 2], self.z+self.size[2]]
		return cuboid


	def _calChildrenAttLogProb(self):
		logp = 0
		if self.children:
			for child in self.children:
				logp += child.calAttLogProb()
			return logp
		else:
			return 0

	def _calChildLogProb(self):
		childLogProb = 0
		if self.children:
			for child in self.node.children:
				childLogProb += pm.distributions.categorical_like(self.childrenCount[child.symbol], self.node.childrenProb[child.symbol])	
		return childLogProb

	def calAttLogProb(self):
		d = self.node.distribution
		self.attLogProb = 0

		# Calculate the log probability of position
		pmix = 0
		for i in range(0, len(d.xPosProbs)):
			logp = pm.distributions.normal_like(self.x, self.xPos[i], 1/d.xVar[i])
			pmix += d.xPosProbs[i] * np.exp(logp)
		pmix += smallNumber
		self.attLogProb += np.log(pmix)

		pmix = 0
		for i in range(0, len(d.yPosProbs)):
			logp = pm.distributions.normal_like(self.y, self.yPos[i], 1/d.yVar[i])
			pmix += d.yPosProbs[i] * np.exp(logp)
		pmix += smallNumber
		self.attLogProb += np.log(pmix)

		pmix = 0
		for i in range(0, len(d.zPosProbs)):
			logp = pm.distributions.normal_like(self.z, self.zPos[i], 1/d.zVar[i])
			pmix += d.zPosProbs[i] * np.exp(logp)
		pmix += smallNumber
		self.attLogProb += np.log(pmix)

		# Calculate the log probability of size
		pmix = 0
		for i in range(0, len(d.sizeProbs)):
			logp = pm.distributions.mv_normal_cov_like(self.size, d.sizeMu[i], np.diag(d.sizeVar[i]))
			pmix += d.sizeProbs[i] * np.exp(logp)
		pmix += smallNumber
		self.attLogProb += np.log(pmix)

		# Calculate the log probability of orientation
		pmix = 0
		for i in range(0, len(d.thetaProbs)):
			logp = pm.distributions.normal_like(self.angleNormalize(self.theta - self.thetaBase - d.thetaMu[i]), 0, 1/d.thetaVar[i])
			pmix += d.thetaProbs[i] * np.exp(logp)
		pmix += smallNumber
		self.attLogProb += np.log(pmix)

		# NORMALIZE
		#return self.attLogProb/7.0
		return self.attLogProb

	def _calLogPrior(self):
		if self.node.isTerminal:
			self.prior = self.calAttLogProb()
		else:
			self.prior = self._calChildLogProb() + self.calAttLogProb()
			for c in self.children:
				self.prior += c._calLogPrior()

		return self.prior
	
	# NORMALIZE
	def calNormalizedLogPrior(self):
		#return self._calLogPrior()/self.entitiesCount()
		return self._calLogPrior()
	

	def entitiesCount(self):
		n = 1
		if self.children:
			for child in self.children:
				n += child.entitiesCount()
		return n

	def _randomEntityRecur(self, n):
		remainingNodes = n - 1

		if np.random.rand() <= 1./n:
			return self
		elif self.children != None:
			for child in self.children:
				randomChild = child._randomEntityRecur(remainingNodes)
				if type(randomChild) is tuple:
					remainingNodes = randomChild[1]
				else:
					return randomChild
			return None, remainingNodes
		else:
			return None, remainingNodes
	
	def randomEntity(self):
		return self._randomEntityRecur(self.entitiesCount())

	def parentEntitiesCount(self):
		n = 0
			
		if self.children:
			n += 1
			for child in self.children:
				n += child.parentEntitiesCount()
		return n
		
	def _randomParentEntityRecur(self, n):
		if n == 0:
			return None
		
# 		if self.children != None:
		if self.children:
			if np.random.rand() <= 1./n:
				return self
			else:
				remainingNodes = n-1
				for child in self.children:
					randomChild = child._randomParentEntityRecur(remainingNodes)
					if type(randomChild) is tuple:
						remainingNodes = randomChild[1]
					else:
						return randomChild
				return None, remainingNodes
		else:
			return None, n
		
	def randomParentEntity(self):
		return self._randomParentEntityRecur(self.parentEntitiesCount())
	
	def termEntitiesCount(self):
		n = 0
		if self.node.isEntity:
			n += 1
			
		if self.children:
			for child in self.children:
				n += child.termEntitiesCount()
		return n
		
	def _randomTermEntityRecur(self, n):
		remainingNodes = n
		if self.node.isEntity:
			remainingNodes -= 1
			if np.random.rand() <= 1./n:
				return self
		
		if self.children:
			for child in self.children:
				randomChild = child._randomTermEntityRecur(remainingNodes)
				if type(randomChild) is tuple:
					remainingNodes = randomChild[1]
				else:
					return randomChild
			return None, remainingNodes
		else:
			return None, remainingNodes
		
	def randomTermEntity(self):
		# Returns an entity that "is entity" -- will be rendered on the image.
		# Notice that the root node will also be considered as "is entity".
		return self._randomTermEntityRecur(self.termEntitiesCount())
	
	def nonTermEntitiesCount(self):
		n = 0
		if not self.node.isTerminal:
			n += 1
			
		if self.children:
			for child in self.children:
				n += child.nonTermEntitiesCount()
		return n
		
	def _randomNonTermEntityRecur(self, n):
		if self.node.isTerminal:
			return None, n
		
		remainingNodes = n - 1
		if np.random.rand() <= 1./n:
			return self
		elif self.children:
			for child in self.children:
				randomChild = child._randomNonTermEntityRecur(remainingNodes)
				if type(randomChild) is tuple:
					remainingNodes = randomChild[1]
				else:
					return randomChild
			return None, remainingNodes
		else:
			return None, remainingNodes
		
	def randomNonTermEntity(self):
		return self._randomNonTermEntityRecur(self.nonTermEntitiesCount())
	
	def proposeFromParent(self):
		self.sampleAttribute()
		coor = self._getWorldCoor()
		cuboid = [coor[0, 0], coor[1, 0], self.z, coor[0, 2], coor[1, 2], self.z+self.size[2]]
		self.assignCuboid(cuboid)
		
	def proposeFromChildLogProb(self):
		totalProb = 0
		
		for child in self.children:
			dParent = self.node.distribution
			dChild = child.node.distribution
			prob = 1
	
			# Calculate the log probability of size
			pmix = 0
			for i in range(0, len(dParent.sizeProbs)):
				logp = pm.distributions.mv_normal_cov_like(self.size, dParent.sizeMu[i], np.diag(dParent.sizeVar[i]))
				pmix += dParent.sizeProbs[i] * np.exp(logp)
			pmix += smallNumber
			prob *= pmix
			
			# Calculate the log probability of position
			pmix = 0
			for i in range(0, len(dChild.xPosProbs)):
				logp = pm.distributions.normal_like(child.x, child.xPos[i], 1/dChild.xVar[i])
				pmix += dChild.xPosProbs[i] * np.exp(logp)
			pmix += smallNumber
			prob *= pmix
	
			pmix = 0
			for i in range(0, len(dChild.yPosProbs)):
				logp = pm.distributions.normal_like(child.y, child.yPos[i], 1/dChild.yVar[i])
				pmix += dChild.yPosProbs[i] * np.exp(logp)
			pmix += smallNumber
			prob *= pmix
	
			pmix = 0
			for i in range(0, len(dChild.zPosProbs)):
				logp = pm.distributions.normal_like(child.z, child.zPos[i], 1/dChild.zVar[i])
				pmix += dChild.zPosProbs[i] * np.exp(logp)
			pmix += smallNumber
			prob *= pmix
	
			# Calculate the log probability of orientation
			pmix = 0
			for i in range(0, len(dChild.thetaProbs)):
				logp = pm.distributions.normal_like(self.angleNormalize(child.theta - child.thetaBase - dChild.thetaMu[i]), 0, 1/dChild.thetaVar[i])
				pmix += dChild.thetaProbs[i] * np.exp(logp)
			pmix += smallNumber
			prob *= pmix
			
			# NORMALIZE
			#totalProb += np.power(prob, 1./7)
			totalProb += prob

		return np.log(totalProb/len(self.children))
	
	def proposeFromChild(self, child):
		def baseOrientation():
			distIdx = np.argsort([abs(y), abs(x), abs(newSize[1] - y), abs(newSize[0] - x)])
			thetaBase = distIdx[0]*np.pi/2
			return thetaBase
		
		if not self.children:
			print "Warning: proposeFromChild called by non-parent"
			print self.parentEntitiesCount()
			print self.childrenCount
			print self.children
			
			return False
		
		dParent = self.node.distribution
		dChild = child.node.distribution
		
		# Sample a new size for the current entity
		newSizeBranch = pm.distributions.rcategorical(dParent.sizeProbs)
		newSize = pm.distributions.rmv_normal_cov(dParent.sizeMu[newSizeBranch], np.diag(dParent.sizeVar[newSizeBranch]))
		
		# Sample the position of the chosen child
		xPos = [0, 0, 0]
		yPos = [0, 0, 0]
		zPos = [0, 0, 0]
		for i in range(3):
			xPos[i] = 0.5 * i * newSize[0] + dChild.xMu[i]
			yPos[i] = 0.5 * i * newSize[1] + dChild.yMu[i]
			zPos[i] = 0.5 * i * newSize[2] + dChild.zMu[i]
			
		xPosBranch = pm.distributions.rcategorical(dChild.xPosProbs)
		yPosBranch = pm.distributions.rcategorical(dChild.yPosProbs)
		# zPosBranch = pm.distributions.rcategorical(dChild.zPosProbs)
		x = pm.distributions.rnormal(xPos[xPosBranch], 1/dChild.xVar[xPosBranch])
		y = pm.distributions.rnormal(yPos[yPosBranch], 1/dChild.yVar[yPosBranch])
		# z = pm.distributions.rnormal(zPos[zPosBranch], 1/dChild.zVar[zPosBranch])
		
		# Sample a new z for the current entity
		newZPosBranch = pm.distributions.rcategorical(dParent.zPosProbs)
		newZ = pm.distributions.rnormal(self.zPos[newZPosBranch], 1/dParent.zVar[newZPosBranch])
		
		# Sample the orientation of the chose child
		thetaBranch = pm.distributions.rcategorical(dChild.thetaProbs)
		childTheta = baseOrientation() + pm.distributions.rnormal(dChild.thetaMu[thetaBranch], 1/dChild.thetaVar[thetaBranch])
		
		
		
		H1 = np.matrix([[1, 0, -child.node.anchor[0]*child.size[0]], [0, 1, -child.node.anchor[1]*child.size[1]], [0, 0, 1]])

		H2 = np.matrix([[np.cos(childTheta), np.sin(childTheta), x], \
					  	[-np.sin(childTheta), np.cos(childTheta), y], \
					  	[0, 0, 1]])
		childToParent = H2 * H1
		
		parentToWorld = child._getWorldCoor()[0:3, 0:3] * np.linalg.inv(childToParent*child._getCorners()[0:3, 0:3])
		coor = parentToWorld * self._getCorners()
		cuboid = [coor[0, 0], coor[1, 0], newZ, coor[0, 2], coor[1, 2], newZ+newSize[2]]
		self.assignCuboid(cuboid)
		
		
	def proposeGroupFromParent(self):
		self.sampleAttribute()
		coor = self._getWorldCoor()
		cuboid = [coor[0, 0], coor[1, 0], self.z, coor[0, 2], coor[1, 2], self.z+self.size[2]]
		self.assignGroupCuboid(cuboid)
		
	def proposeGroupFromChildLogProb(self, child):
		dParent = self.node.distribution
		dChild = child.node.distribution
		prob = 1

		# Calculate the log probability of size
		pmix = 0
		for i in range(0, len(dParent.sizeProbs)):
			logp = pm.distributions.mv_normal_cov_like(self.size, dParent.sizeMu[i], np.diag(dParent.sizeVar[i]))
			pmix += dParent.sizeProbs[i] * np.exp(logp)
		pmix += smallNumber
		prob *= pmix
		
		# Calculate the log probability of position
		pmix = 0
		for i in range(0, len(dChild.xPosProbs)):
			logp = pm.distributions.normal_like(child.x, child.xPos[i], 1/dChild.xVar[i])
			pmix += dChild.xPosProbs[i] * np.exp(logp)
		pmix += smallNumber
		prob *= pmix

		pmix = 0
		for i in range(0, len(dChild.yPosProbs)):
			logp = pm.distributions.normal_like(child.y, child.yPos[i], 1/dChild.yVar[i])
			pmix += dChild.yPosProbs[i] * np.exp(logp)
		pmix += smallNumber
		prob *= pmix

		pmix = 0
		for i in range(0, len(dChild.zPosProbs)):
			logp = pm.distributions.normal_like(child.z, child.zPos[i], 1/dChild.zVar[i])
			pmix += dChild.zPosProbs[i] * np.exp(logp)
		pmix += smallNumber
		prob *= pmix

		# Calculate the log probability of orientation
		pmix = 0
		for i in range(0, len(dChild.thetaProbs)):
			logp = pm.distributions.normal_like(self.angleNormalize(child.theta - child.thetaBase - dChild.thetaMu[i]), 0, 1/dChild.thetaVar[i])
			pmix += dChild.thetaProbs[i] * np.exp(logp)
		pmix += smallNumber
		prob *= pmix
		
		# NORMALIZE
		#prob = np.power(prob, 1./7)

		return np.log(prob)
	
	def proposeGroupFromChild(self, child):
		def baseOrientation():
			distIdx = np.argsort([abs(y), abs(x), abs(newSize[1] - y), abs(newSize[0] - x)])
			thetaBase = distIdx[0]*np.pi/2
			return thetaBase
		
		dParent = self.node.distribution
		dChild = child.node.distribution
		
		# Sample a new size for the current entity
		newSizeBranch = pm.distributions.rcategorical(dParent.sizeProbs)
		newSize = pm.distributions.rmv_normal_cov(dParent.sizeMu[newSizeBranch], np.diag(dParent.sizeVar[newSizeBranch]))
		
		# Sample the position of the chosen child
		xPos = [0, 0, 0]
		yPos = [0, 0, 0]
		zPos = [0, 0, 0]
		for i in range(3):
			xPos[i] = 0.5 * i * newSize[0] + dChild.xMu[i]
			yPos[i] = 0.5 * i * newSize[1] + dChild.yMu[i]
			zPos[i] = 0.5 * i * newSize[2] + dChild.zMu[i]
			
		xPosBranch = pm.distributions.rcategorical(dChild.xPosProbs)
		yPosBranch = pm.distributions.rcategorical(dChild.yPosProbs)
		# zPosBranch = pm.distributions.rcategorical(dChild.zPosProbs)
		x = pm.distributions.rnormal(xPos[xPosBranch], 1/dChild.xVar[xPosBranch])
		y = pm.distributions.rnormal(yPos[yPosBranch], 1/dChild.yVar[yPosBranch])
		# z = pm.distributions.rnormal(zPos[zPosBranch], 1/dChild.zVar[zPosBranch])
		
		# Sample a new z for the current entity
		newZPosBranch = pm.distributions.rcategorical(dParent.zPosProbs)
		newZ = pm.distributions.rnormal(self.zPos[newZPosBranch], 1/dParent.zVar[newZPosBranch])
		
		# Sample the orientation of the chose child
		thetaBranch = pm.distributions.rcategorical(dChild.thetaProbs)
		childTheta = baseOrientation() + pm.distributions.rnormal(dChild.thetaMu[thetaBranch], 1/dChild.thetaVar[thetaBranch])
		
		
		
		H1 = np.matrix([[1, 0, -child.node.anchor[0]*child.size[0]], [0, 1, -child.node.anchor[1]*child.size[1]], [0, 0, 1]])

		H2 = np.matrix([[np.cos(childTheta), np.sin(childTheta), x], \
					  	[-np.sin(childTheta), np.cos(childTheta), y], \
					  	[0, 0, 1]])
		childToParent = H2 * H1
		
		parentToWorld = child._getWorldCoor()[0:3, 0:3] * np.linalg.inv(childToParent*child._getCorners()[0:3, 0:3])
		coor = parentToWorld * self._getCorners()
		cuboid = [coor[0, 0], coor[1, 0], newZ, coor[0, 2], coor[1, 2], newZ+newSize[2]]
		self.assignGroupCuboid(cuboid, child)
		

	def delEntity(self):
		if self.parent:
			self.parent.children.remove(self)
			self.parent.childrenCount[self.node.symbol] -= 1
		else:
			print "Error: trying to delete the root entity"
			exit(1)


	def angleNormalize(self, theta):
		theta = theta%(2*np.pi)
		if theta > np.pi:
			return theta - 2*np.pi
		else:
			return theta
		
	def getVol(self):
		return self.length*self.width*self.height

	def getData(self):
		cuboids = self._getWorldCoor()
		cuboids = cuboids[0:2, :].T.flatten().tolist()
		if self.node.isEntity == 0 or self.node.isEntity == 2:
			cuboids[0].extend([self.z, self.z + self.size[2], self.node.isEntity])
		elif self.node.isEntity == 1:
			cuboids[0].extend([self.z, self.z + self.size[2], self.node.isEntity])
			cuboids[0].extend(self.node.color)

		if self.children:
			for child in self.children:
				cuboids.extend(child.getData())

		return cuboids

	def get_entities(self):
		entities = [self]
		if self.children:
			for child in self.children:
				entities.extend(child.get_entities())
		return entities

	def writeToFile(self, outputFile):
		coor = self._getWorldCoor()
		coor = coor[0:2, :].T.flatten()
		if self.node.isEntity == 0 or self.node.isEntity == 2:
			line = '{:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} '.format(\
					coor[0, 0], coor[0, 1], coor[0, 2], coor[0, 3], coor[0, 4], coor[0, 5], coor[0, 6], coor[0, 7], \
					self.z, self.z + self.size[2], self.node.isEntity) + self.node.symbol + '\n'
		elif self.node.isEntity == 1:
			line = '{:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:d} {:d} {:d} '.format(\
					coor[0, 0], coor[0, 1], coor[0, 2], coor[0, 3], coor[0, 4], coor[0, 5], coor[0, 6], coor[0, 7], \
					self.z, self.z + self.size[2], self.node.isEntity, self.node.color[0], self.node.color[1], self.node.color[2]) + self.node.symbol + '\n'

		outputFile.write(line)
		
		if self.children:
			for child in self.children:
				child.writeToFile(outputFile)

